using Crux.Core.LangChain.Abstractions.Constants;
using Crux.Core.LangChain.Abstractions.Databases;
using Crux.Core.LangChain.Abstractions.Databases.Metadatas;
using MongoDB.Bson;
using MongoDB.Driver;
namespace Crux.Core.LangChain.Databases.Mongo;
/// <summary>
/// 
/// </summary>
public class MongoRagEngine : IVectorDatabase, IDisposable
{
    private const string DefaultSchema = "langchain";
    private readonly string FunctionName = DistanceFunction.EuclideanDistance;
    private readonly IMongoClient mongoClient;
    private readonly IMongoDatabase mongoDatabase;
    private bool _disposed = false;
    /// <summary>
    /// 
    /// </summary>
    /// <param name="setting"></param>
    /// <param name="databaseName"></param>
    /// <param name="function"></param>
    /// <exception cref="ArgumentNullException"></exception>
    public MongoRagEngine(
        MongoClientSettings setting,
        string databaseName = DefaultSchema,
        string function = DistanceFunction.EuclideanDistance)
    {
        this.mongoClient = new MongoClient(setting ?? throw new ArgumentNullException(nameof(setting)));
        this.mongoDatabase = mongoClient.GetDatabase(databaseName ?? throw new ArgumentNullException(nameof(databaseName)));
        this.FunctionName = function ?? throw new ArgumentNullException(nameof(function));
    }
    /// <summary>
    /// 
    /// </summary>
    /// <param name="collectionName"></param>
    /// <param name="dimensions"></param>
    /// <param name="cancellationToken"></param>
    /// <returns></returns>
    public async Task<IVectorCollection> GetCollectionAsync(string collectionName, int dimensions = 1536, CancellationToken cancellationToken = default)
    {
        if (this._disposed)
        {
            throw new ObjectDisposedException(nameof(MongoRagEngine));
        }

        if (!await IsCollectionExistsAsync(collectionName, cancellationToken).ConfigureAwait(false))
        {
            await CreateCollectionAsync(collectionName, dimensions, cancellationToken).ConfigureAwait(false);
        }

        return new MongoRagCollection(mongoDatabase, "idx_" + collectionName, name: collectionName);
    }

    /// <inheritdoc />
    public async Task<bool> IsCollectionExistsAsync(string collectionName, CancellationToken cancellationToken = default)
    {
        if (this._disposed)
        {
            throw new ObjectDisposedException(nameof(MongoRagEngine));
        }

        var filter = new BsonDocument("name", collectionName);
        var options = new ListCollectionNamesOptions { Filter = filter };

        var collections = await mongoDatabase.ListCollectionNamesAsync(options, cancellationToken).ConfigureAwait(false);

        return await collections.AnyAsync(cancellationToken).ConfigureAwait(false);
    }
    /// <inheritdoc />
    private async Task CreateCollectionAsync(string collectionName, int dimensions = 1536, CancellationToken cancellationToken = default)
    {
        if (this._disposed)
        {
            throw new ObjectDisposedException(nameof(MongoRagEngine));
        }

        await mongoDatabase.CreateCollectionAsync(collectionName, new CreateCollectionOptions
        {
        }, cancellationToken).ConfigureAwait(false);

        var collection = mongoDatabase.GetCollection<Vector>(collectionName);

        var indexCursor = await collection.Indexes.ListAsync(cancellationToken).ConfigureAwait(false);

        var indexes = indexCursor.ToList(cancellationToken).Select(index => index["name"].ToString()) ?? [];

        var indexName = "idx_" + collectionName;
        if (indexes.Contains(indexName))
        {
            return;
        }

        var fieldsArray = new BsonArray
        {
            new BsonDocument
            {
                { "type", "vector" },
                { "numDimensions", dimensions },
                { "path", nameof(Vector.Embedding) },
                { "similarity", GetDistanceFunction() },
            }
        };

        var indexArray = new BsonArray
            {
                new BsonDocument
                {
                    { "name", indexName },
                    { "type", "vectorSearch" },
                    { "definition", new BsonDocument { ["fields"] = fieldsArray } },
                }
            };

        var createIndexCommand = new BsonDocument
            {
                { "createSearchIndexes", collectionName },
                { "indexes", indexArray }
            };

        await this.mongoDatabase.RunCommandAsync<BsonDocument>(createIndexCommand, cancellationToken: cancellationToken).ConfigureAwait(false);
    }

    private string GetDistanceFunction()
    {
        return FunctionName switch
        {
            DistanceFunction.CosineSimilarity => "cosine",
            DistanceFunction.DotProductSimilarity => "dotProduct",
            DistanceFunction.EuclideanDistance => "euclidean",
            _ => throw new InvalidOperationException($"Distance function '{FunctionName}'is not supported by the MongoDB VectorStore.")
        };
    }
    /// <inheritdoc />
    public async Task DeleteCollectionAsync(string collectionName, CancellationToken cancellationToken = default)
    {
        if (this._disposed)
        {
            throw new ObjectDisposedException(nameof(MongoRagEngine));
        }

        await mongoDatabase.DropCollectionAsync(collectionName, cancellationToken).ConfigureAwait(false);
    }
    /// <summary>
    /// 
    /// </summary>
    /// <param name="options"></param>
    /// <param name="cancellationToken"></param>
    /// <returns></returns>
    public async Task<IClientSessionHandle> StartSessionAsync(
        ClientSessionOptions? options = null,
        CancellationToken cancellationToken = default)
    {
        return await mongoClient.StartSessionAsync(options, cancellationToken).ConfigureAwait(false);
    }
    /// <summary>
    /// 
    /// </summary>
    public void Dispose()
    {
        if (_disposed)
        {
            return;
        }

        _disposed = true;
    }
}